#include "log.h"
#include "dial.h"
#include "conn_track.h"

#define IS_RESPONSE(flag) ((flag) >> 15)
#define REPLY_CDOE(flag) ((flag)&0xf)

static char* Dns_Error_String[] = {
    "No Error",
    "Format Error",
    "Server Failure",
    "Non-Existent Domain",
    "Not Implemented",
    "Query Refused",
    "Name Exists when it should not",
    "RR Set Exists when it should not",
    "RR Set that should exist does not",
    "Server Not Authoritative for zone",
    "Not Authorized",
    "Name not contained in zone",
    "Bad OPT Version",
    "TSIG Signature Failure",
    "Key not recognized",
    "Signature out of time window",
    "Bad TKEY Mode",
    "Duplicate key name",
    "Algorithm not supported",
    "Bad Truncation",
    "Bad/missing Server Cookie",
};

// void* _do_bridge_dns(dial_arg_t* dial_arg, dial_result_t* dial_result);

void dns_result_print(dns_result_t* result)
{
    int i = 0;
    struct in_addr addr;

    log_trace("latency:%.2f ms", result->latency);
    log_trace("rcode:%d (%s)", result->rcode, Dns_Error_String[result->rcode]);
    for (i = 0; i < sizeof(result->ipaddr) / sizeof(uint32_t); i++) {
        addr.s_addr = result->ipaddr[i];
        log_trace("ipaddr:%s", inet_ntoa(addr));
    }
}

// void* do_dns(void* arg)
// {
//     int ret;
//     struct timeval s_tm, e_tm;
//     double diff_time;
//     struct addrinfo *answer, *cur;
//     dial_arg_t* dial_arg = (dial_arg_t*)arg;
//     dial_result_t* dial_result = calloc(1, sizeof(dial_result_t));
//     dns_arg_t* dns_arg = &dial_arg->in.dns_arg;
//     dns_result_t* dns_result = &dial_result->out.dns_result;

//     if (dial_arg->bridge_mode) {
//         _do_bridge_dns(dial_arg, dial_result);
//     } else {
//         gettimeofday(&s_tm, NULL);
//         answer = wd_gethostbyname(dns_arg->domain, dns_result->ipaddr, sizeof(dns_result->ipaddr) / sizeof(uint32_t));
//         gettimeofday(&e_tm, NULL);
//         freeaddrinfo(answer);
//         diff_time = TV2MS(&e_tm) - TV2MS(&s_tm);
//         dns_result->latency = diff_time;
//     }

//     return (void*)dial_result;
// }

// /////////////////////// dns bridge mode packet ///////////

#define DNS_QUERY_LENGTH(domain) (1 + strlen(domain) + 1 + 2 * sizeof(uint16_t))
dns_t* dns_packet_new(int query_len, int* len)
{
    int total_len = sizeof(dns_t) + query_len;
    dns_t* dp = malloc(total_len);
    if (dp == NULL) {
        log_error("New dns packet failed.");
        return NULL;
    }
    memset(dp, 0, total_len);
    *len = total_len;
    return dp;
}
static int gen_dns_query(char* data, char* domain_name)
{
    char buf[256] = { 0 };
    char* delim = ".";
    char *p, *q;
    int len = 0;

    strcpy(buf, domain_name);
    q = data;
    p = strtok(buf, delim);
    len = strlen(p);
    *(q++) = len;
    strcpy(q, p);
    q += len;

    while ((p = strtok(NULL, delim))) {
        len = strlen(p);
        *(q++) = len;
        strcpy(q, p);
        q += len;
    }

    // log_buf(data, strlen(domain_name) + 2);
    return strlen(domain_name) + 2;
}
dns_t* dns_packet(char* domain, int* len)
{
    int query_name_len;
    uint16_t *query_type, *query_class;
    dns_t* dp;

    dp = dns_packet_new(DNS_QUERY_LENGTH(domain), len);
    if (!dp) {
        return NULL;
    }
    dp->id = (uint16_t)get_random();
    dp->flags = htons(0x0100);
    dp->questions = htons(0x00001);
    query_name_len = gen_dns_query(dp->data, domain);
    query_type = (uint16_t*)(dp->data + query_name_len);
    *query_type = htons(0x0001);
    query_class = query_type + 1;
    *query_class = htons(0x0001);
    return dp;
}
void dns_packet_free(dns_t* dp)
{
    free(dp);
}

typedef struct {
    union {
        uint16_t ptr;
        char name[128];
    } name;
    uint16_t type;
    uint16_t class;
    uint32_t ttl;
    uint16_t d_len;
} _answer_t;

static int get_query_len(char* buf)
{
#define _type uint16_t
#define _class uint16_t

    char* ptr = NULL;
    int query_len = 0;

    ptr = buf;
    while (*ptr) {
        query_len += *ptr;
        query_len += 1;
        ptr += *ptr + 1;
    }
    query_len += 1;
    query_len += sizeof(_type);
    query_len += sizeof(_class);
    return query_len;
}
static void get_addrs(uint8_t* answer, uint32_t* addrs, int addr_size)
{
    int i = 0;
    _answer_t a;
    uint8_t* ptr = NULL;
    ptr = answer;
    while (*ptr) {
        if (*ptr == 0xc0) {
            a.name.ptr = *(uint16_t*)ptr;
            ptr += 2;

            a.type = *(uint16_t*)ptr;
            ptr += 2;

            a.class = *(uint16_t*)ptr;
            ptr += 2;

            a.ttl = *(uint16_t*)ptr;
            ptr += 4;

            a.d_len = *(uint16_t*)ptr;
            ptr += 2;

            if (a.type == htons(0x5)) {
            } else if (a.type == htons(0x1)) {
                addrs[i++] = *(uint32_t*)ptr;
            }
            ptr += ntohs(a.d_len);
        } else {
            log_debug("dns answer first byte not 0xc0");
            log_error("dns answer error, open log debug see detail.");
        }
    }
}

static dial_result_t* _do_getaddrinfo(dns_arg_t* arg, dial_result_t* dial_result)
{
    int ret, i = 0;
    struct timeval s_tm, e_tm;
    struct addrinfo *answer, hint, *cur;
    dns_result_t* result = &dial_result->out.dns_result;

    bzero(&hint, sizeof(hint));
    hint.ai_family = AF_INET;
    hint.ai_socktype = SOCK_STREAM;

    gettimeofday(&s_tm, NULL);
    if ((ret = getaddrinfo(arg->domain, NULL, &hint, &answer)) != 0) {
        log_error("_do_getaddrinfo error(%d):%s", ret, gai_strerror(ret));
        dial_result->err = DNS_ERR_GETADDR;
        return dial_result;
    }
    gettimeofday(&e_tm, NULL);

    for (cur = answer; cur; cur = cur->ai_next) {
        result->ipaddr[i++] = ((struct sockaddr_in*)cur->ai_addr)->sin_addr.s_addr;
        if (i < Array_Size(result->ipaddr))
            continue;
        else
            break;
    }
    result->rcode = ret;
    result->latency = TV2MS(&e_tm) - TV2MS(&s_tm);
    freeaddrinfo(answer);
    return dial_result;
}

void* do_dns(void* arg)
{
    int DNS_fd = 0;
    char buf[1500] = { 0 }, *p_query = NULL, *p_answer = NULL;
    int ret;
    dns_t *dns, *dns_resp;
    int dns_len = 0, query_len = 0;
    struct sockaddr_in svr_addr;
    struct in_addr addr;
    uint16_t svr_port;
    dial_arg_t* dial_arg = (dial_arg_t*)arg;
    dial_result_t* dial_result = calloc(1, sizeof(dial_result_t));
    dns_arg_t* dns_arg = &dial_arg->in.dns_arg;
    dns_result_t* dns_result = &dial_result->out.dns_result;
    struct timeval s_tm, e_tm;

    #ifdef BRIDGE_MODE_SUPPORT
    conn_track ct = { 0 };
    if (dial_arg->bridge_mode) {

        ct.dest_ip = conn_dns_server();
        ct.dest_port = htons(53);
        ct.proto = IPPROTO_UDP;

        DNS_fd = conn_track_regist(&ct);
// log_trace("DNS_fd:%d", DNS_fd);
    } else 
    #endif
    {
        addr.s_addr = inet_addr(dns_arg->dns_svr);
        if (addr.s_addr == 0) {
            return _do_getaddrinfo(dns_arg, dial_result);
        }
        DNS_fd = socket_dgram_new();
        sock_rcvtmo(DNS_fd, 3, 0);
        svr_port = dns_arg->port ? dns_arg->port : 53;
        sockaddr_set(&svr_addr, AF_INET, htons(svr_port), dns_arg->dns_svr);
        connect(DNS_fd, (struct sockaddr*)&svr_addr, sizeof(svr_addr));
    }

    dns = dns_packet(dns_arg->domain, &dns_len);
    gettimeofday(&s_tm, 0);
    if (0 > write(DNS_fd, dns, dns_len)) {
        log_syserr("do dns write error.");
        dial_result->err = DNS_ERR_WRITE;
        goto out;
    }
    ret = read(DNS_fd, buf, sizeof(buf));
    if (ret < 0) {
        log_syserr("do dns read error.");
        dial_result->err = DNS_ERR_READ;
        goto out;
    }
    log_buf(LOG_DEBUG, buf, ret);
    gettimeofday(&e_tm, 0);
    dns_resp = (dns_t*)buf;
    dns_result->rcode = REPLY_CDOE(ntohs(dns_resp->flags));
    if (dns_result->rcode == 0) {
        // log_buf(LOG_DEBUG, buf, ret);
        // log_debug("answers:%x", ntohs(dns_resp->answer_RRs));
        // log_debug("IS_RESPONSE:%x", IS_RESPONSE(ntohs(dns_resp->flags)));
        p_query = buf + sizeof(dns_t);
        query_len = get_query_len(p_query);
        // log_debug("query_len:%d", query_len);
        p_answer = p_query + query_len;
        // log_buf(LOG_DEBUG, p_answer, 64);
        get_addrs((uint8_t*)p_answer, dns_result->ipaddr, sizeof(dns_result->ipaddr));
        // log_buf(LOG_DEBUG, dns_result->ipaddr, sizeof(dns_result->ipaddr));
        dns_result->latency = TV2MS(&e_tm) - TV2MS(&s_tm);
    }

out:
    close(DNS_fd);
    dns_packet_free(dns);
    return (void*)dial_result;
}
